#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <cluster.h>
#include <matrix.h>

#include <vector>
#include <iostream>
#include <iomanip>
#include <algorithm>

using namespace sds;
using namespace std;

int main(int argc, const char* argv[]) {
  SFile reffile(argv[2]);
  SDataView<unsigned char> ref =
    reffile.getDataset(0).getDataView<unsigned char>();

  vector<SInteger> size = ref.getSize();

  SFile srcfile(argv[1]);
  SDataView<Point> src(size);

  for (int k = 0; k < CH; ++k) {
    SDataView<unsigned char> view =
      srcfile.getDataset(k).getDataView<unsigned char>();
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	((Point&)src[i][j])[k] = (double)view[i][j];
      }    
    }
  }

  SDataView<unsigned char> trg(size);
  SDataView<bool> rd(size);

  vector<Cluster> clusters(256);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ref[i][j] != 0) {
	clusters[ref[i][j]] += src[i][j];
	trg[i][j] = ref[i][j];
	rd[i][j] = true;
      } else {
	rd[i][j] = false;
      }
    }
  }

  for (int s = 16; s > 0; --s) {
    cout << "Step: " << s << endl;

    vector<int> index;
    vector<Vector> expected;
    vector<Matrix> root;

    for (int k = 1; k < 256; ++k) {
      if (clusters[k].size() == 0) continue;
      index.push_back(k);
      Vector exp(CH);
      for (int i = 0; i < CH; ++i) {
	exp(i) = clusters[k].expected(i);
      }
      expected.push_back(exp);
      Matrix cov(CH, CH);
      for (int i = 0; i < CH; ++i) {
	for (int j = 0; j < CH; ++j) {
	  cov(i, j) = clusters[k].covariance(i, j);
	}
      }
      root.push_back(cov.inverse().cholesky());
    }

    vector<pair<double, int> > diff;
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	if (rd[i][j]) continue;
	Vector current(CH);
	for (int l = 0; l < CH; ++l) {
	  current(l) = ((Point&)src[i][j])[l];
	}
	int minPos = -1; 
	double minVal = 0.0;
	for (int k = 0; k < index.size(); ++k) {
	  Vector norm = root[k] * (current - expected[k]);
	  double dist = 0.0;
	  for (int l = 0; l < CH; ++l) {
	    dist += norm(l) * norm(l);
	  }
	  if (minPos == -1 || minVal > dist) {
	    minPos = k;
	    minVal = dist;
	  }
	}
	trg[i][j] = index[minPos];
	diff.push_back(make_pair(minVal, i * size[1] + j));
      }

      vector<pair<double, int> >::iterator 
	it, lt = diff.begin() + diff.size() / s;
      nth_element(diff.begin(), lt, diff.end());

      for (it = diff.begin(); it != lt; ++it) {
	rd[it->second / size[1]][it->second % size[1]] = true;
	clusters[trg[it->second / size[1]][it->second % size[1]]] 
	  += src[it->second / size[1]][it->second % size[1]];	
      }

    }
  }
  
  SFile trgfile(argv[3], SFile::CREATE);
  trgfile.createDataset("clasified", UINT8, size).
    setDataView<unsigned char>(trg);
  return 0;
}
